package com.zhou.util;

import com.google.common.collect.Lists;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.http.HttpHost;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.MatchQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.net.UnknownHostException;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class ESSource extends RichParallelSourceFunction<Map<String,String>> {

    private static RestHighLevelClient client=null;

    private volatile boolean isRunning = true;


    public static RestHighLevelClient initClient(ParameterTool parameterTool) throws UnknownHostException {
         client = new RestHighLevelClient(
                RestClient.builder(new HttpHost(parameterTool.get(PropertiesKey.ES_IP), parameterTool.getInt(PropertiesKey.ES_PORT), "http")
                        ));
        return client;
    }

    @Override
    public void run(SourceContext sourceContext) throws Exception {
        SearchRequest searchRequest = new SearchRequest("community_im_chat");
        searchRequest.types("_doc");
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.from(0);
        sourceBuilder.size(10000);
        MatchQueryBuilder chatTypeBuilder = QueryBuilders.matchQuery("chat_type", "chat");
        RangeQueryBuilder timestamp = QueryBuilders.rangeQuery("timestamp");
        timestamp.gte(1588694400000l);
        RangeQueryBuilder lt = timestamp.lt(1588780800000l);
        BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
        boolQuery.must().addAll(Lists.newArrayList(chatTypeBuilder, timestamp, QueryBuilders.existsQuery("chat_type")));
        sourceBuilder.timeout(new TimeValue(60, TimeUnit.SECONDS));
        sourceBuilder.query(boolQuery);
        searchRequest.source(sourceBuilder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        SearchHits hits = searchResponse.getHits();
        long totalHits = hits.getTotalHits().value;
        SearchHit[] searchHits = hits.getHits();
        for (SearchHit hit : searchHits) {
            sourceContext.collect(hit.getSourceAsMap());
            if (isRunning){
                return;
            }
        }
    }

    @Override
    public void cancel() {
        isRunning=false;
    }

}
